#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@his-pc>
#
# Distributed under terms of the MIT license.

"""
    基于kruskal算法的最小生成树.
"""

import heapq


class Kruskal:
    '''
    kruskal算法针对无向图.
    这里假设输入是一颗连通图.
    '''

    def __init__(self, edges, vertex_nums):
        # edges: (weight, start, end)   # 因为要针对weight排序, heapq不支持key比较参数, 所以weight放在最前面
        # vertex编号从1开始
        heapq.heapify(edges)
        self.edges = edges
        self.vertex_nums = vertex_nums
        self.disjoint_parent = {i:-1 for i in range(1, vertex_nums + 1)}        # make union_by_rank set
        self.tree_edge_cnt = 0                                                  # 也可以换成已经在树中节点个数
        self._minimal_edges = []

        while not self.span_over and self.edges:
            weight, start, end = cur_edge = heapq.heappop(self.edges)
            if not self.union_by_rank(start, end):
                continue

            self.tree_edge_cnt += 1
            self._minimal_edges.append(cur_edge)

        # self._minimal_edges.sort(key=lambda d: d[1])
        # if self.span_over:
        #     return self._minimal_edges
        # else:
        #     return []

    @property
    def span_over(self):
        # 是否匹配结束
        return self.tree_edge_cnt == self.vertex_nums - 1

    @property
    def minimal_edges(self):
        # 是否匹配结束
        if self.span_over:
            return self._minimal_edges
        else:
            return []

    def find_parent(self, node_label):
        # 注意: 这里递归压缩parent路径
        parent_label = self.disjoint_parent[node_label]
        if parent_label < 0:
            return node_label

        self.disjoint_parent[node_label] = self.find_parent(parent_label)
        return self.disjoint_parent[node_label]

    def union_by_rank(self, n1, n2):
        # union by depth/rank
        p1 = self.find_parent(n1)
        p2 = self.find_parent(n2)


        if p1 == p2:
            return False

        rank_1 = -self.disjoint_parent[p1]
        rank_2 = -self.disjoint_parent[p2]

        if rank_1 <= rank_2:
            little_rank = rank_1
            great_rank = rank_2
            child = p1
            parent = p2
        else:
            little_rank = rank_2
            great_rank = rank_1
            child = p2
            parent = p1

        new_rank = max(little_rank + 1, great_rank)
        self.disjoint_parent[child] = parent
        self.disjoint_parent[parent] = -new_rank

        return True


def test_kruskal():
    # 见图 ./kruskal_example_1.gif
    # 见图 ./kruskal_example_2.jpg

    from kruskal_data import *

    k = Kruskal(edges_example_1, 6)
    print(k.minimal_edges)


    k = Kruskal(edges_example_2, 10)
    print(k.minimal_edges)


def main():
    print("start main")
    test_kruskal()


if __name__ == "__main__":
    main()
